#!/usr/bin/env python
# coding:utf-8
from PoboAPI import *
import datetime
import numpy as np
import pandas as pd
from collections import deque

# 高频交易策略示例 - 基于Tick级别的价格压力指标(PPI)
# 策略逻辑：通过分析Tick数据中的买卖盘口压力，进行短线交易

def OnStart(context):
    print("高频交易策略启动...")
    
    # 登录交易账号
    context.myacc = None
    if "回测期货" in context.accounts:
        print("登录交易账号[回测期货]")
        if context.accounts["回测期货"].Login():
            context.myacc = context.accounts["回测期货"]
    
    # 初始化参数
    g.exchange = 'SHFE'  # 上海期货交易所
    g.product = 'au'     # 黄金期货
    g.code = GetMainContract(g.exchange, g.product, 20)  # 获取主力合约
    print(f"交易品种: {g.code}")
    
    # 策略参数
    g.tick_window = 50        # Tick数据窗口大小
    g.tick_buffer = deque(maxlen=g.tick_window)  # 用于存储Tick数据的循环队列
    g.threshold_open = 0.7    # 开仓阈值
    g.threshold_close = 0.3   # 平仓阈值
    g.position = 0            # 持仓状态：0-空仓，1-多头，-1-空头
    g.max_position = 3        # 最大持仓手数
    g.stop_loss_pct = 0.002   # 止损比例
    g.take_profit_pct = 0.004 # 止盈比例
    g.entry_price = 0         # 开仓价格
    g.trade_count = 0         # 交易次数
    g.win_count = 0           # 盈利次数
    
    # 订阅行情
    SubscribeQuote(g.code)
    
    # 设置收盘前平仓闹钟
    g.alarm_id_day = SetAlarm(datetime.time(14, 55), RepeatType.Daily)
    g.alarm_id_night = SetAlarm(datetime.time(2, 25), RepeatType.Daily)
    
    # 初始化日志
    g.trade_log = []

# 行情初始化事件
def OnMarketQuotationInitialEx(context, exchange, daynight):
    if exchange != g.exchange:
        return
    
    # 检查主力合约是否变更
    new_code = GetMainContract(g.exchange, g.product, 20)
    if g.code != new_code:
        print(f"主力合约变更: {g.code} -> {new_code}")
        
        # 平掉旧合约持仓
        if g.position != 0 and context.myacc is not None:
            close_all_positions(context)
        
        # 取消旧合约订阅
        UnsubscribeQuote(g.code)
        
        # 更新合约代码
        g.code = new_code
        
        # 订阅新合约
        SubscribeQuote(g.code)
        
        # 重置数据
        g.tick_buffer.clear()
        g.position = 0
        g.entry_price = 0

# 闹钟事件
def OnAlarm(context, alarmid):
    # 收盘前平仓
    if g.position != 0 and context.myacc is not None:
        print("收盘前平仓")
        close_all_positions(context)
        
        # 输出交易统计
        win_rate = g.win_count / g.trade_count if g.trade_count > 0 else 0
        print(f"交易统计: 总交易次数 {g.trade_count}, 盈利次数 {g.win_count}, 胜率 {win_rate:.2%}")
        
        # 记录交易日志
        if g.trade_log:
            log_df = pd.DataFrame(g.trade_log)
            print(log_df.tail())

# 实时行情事件
def OnQuote(context, code):
    if code != g.code:
        return
    
    # 获取最新行情
    quote = GetQuote(code)
    
    # 将Tick数据添加到缓冲区
    tick_data = {
        'time': GetCurrentTime(),
        'price': quote.now,
        'volume': quote.volume,
        'ask_price': quote.askprice,
        'ask_volume': quote.askvolume,
        'bid_price': quote.bidprice,
        'bid_volume': quote.bidvolume
    }
    g.tick_buffer.append(tick_data)
    
    # 当积累了足够的Tick数据后开始计算指标
    if len(g.tick_buffer) >= g.tick_window:
        # 计算价格压力指标(PPI)
        ppi = calculate_ppi()
        
        # 当前价格
        current_price = tick_data['price']
        
        # 交易逻辑
        if context.myacc is None:
            print("交易账户未登录，无法交易")
            return
        
        # 检查止盈止损
        if g.position != 0:
            check_stop_loss_take_profit(context, current_price)
        
        # 开仓信号
        if g.position == 0:  # 空仓状态
            if ppi > g.threshold_open:  # 买方压力大，做多
                open_long_position(context, current_price)
            elif ppi < -g.threshold_open:  # 卖方压力大，做空
                open_short_position(context, current_price)
        
        # 平仓信号
        elif g.position > 0:  # 多头持仓
            if ppi < -g.threshold_close:  # 卖方压力增加，平多
                close_long_position(context, current_price)
        
        elif g.position < 0:  # 空头持仓
            if ppi > g.threshold_close:  # 买方压力增加，平空
                close_short_position(context, current_price)

# 计算价格压力指标(PPI)
def calculate_ppi():
    if len(g.tick_buffer) < g.tick_window:
        return 0
    
    # 提取买卖盘口数据
    bid_volumes = [tick['bid_volume'] for tick in g.tick_buffer]
    ask_volumes = [tick['ask_volume'] for tick in g.tick_buffer]
    
    # 计算买卖盘口压力比
    total_bid_volume = sum(bid_volumes)
    total_ask_volume = sum(ask_volumes)
    
    if total_bid_volume + total_ask_volume == 0:
        return 0
    
    # 计算PPI指标，范围[-1, 1]
    ppi = (total_bid_volume - total_ask_volume) / (total_bid_volume + total_ask_volume)
    
    return ppi

# 开多头仓位
def open_long_position(context, price):
    if g.position != 0:
        return
    
    volume = g.max_position
    price_type = PriceType(PbPriceType.Limit, 16, 0)
    
    # 检查可用资金
    context.myacc.RefreshMargin(g.code, BSType.BuyOpen)
    available_volume = context.myacc.GetValidVolume(g.code, BSType.BuyOpen, price)
    
    if available_volume < volume:
        volume = available_volume
    
    if volume <= 0:
        print("资金不足，无法开仓")
        return
    
    # 执行开仓
    QuickInsertOrder(context.myacc, g.code, 'buy', 'open', price_type, volume)
    g.position = volume
    g.entry_price = price
    
    print(f"开多头仓位: {g.code}, 价格: {price}, 数量: {volume}")
    
    # 记录交易
    g.trade_log.append({
        'time': GetCurrentTime(),
        'action': 'open_long',
        'code': g.code,
        'price': price,
        'volume': volume
    })

# 开空头仓位
def open_short_position(context, price):
    if g.position != 0:
        return
    
    volume = g.max_position
    price_type = PriceType(PbPriceType.Limit, 16, 0)
    
    # 检查可用资金
    context.myacc.RefreshMargin(g.code, BSType.SellOpen)
    available_volume = context.myacc.GetValidVolume(g.code, BSType.SellOpen, price)
    
    if available_volume < volume:
        volume = available_volume
    
    if volume <= 0:
        print("资金不足，无法开仓")
        return
    
    # 执行开仓
    QuickInsertOrder(context.myacc, g.code, 'sell', 'open', price_type, volume)
    g.position = -volume
    g.entry_price = price
    
    print(f"开空头仓位: {g.code}, 价格: {price}, 数量: {volume}")
    
    # 记录交易
    g.trade_log.append({
        'time': GetCurrentTime(),
        'action': 'open_short',
        'code': g.code,
        'price': price,
        'volume': volume
    })

# 平多头仓位
def close_long_position(context, price):
    if g.position <= 0:
        return
    
    volume = abs(g.position)
    price_type = PriceType(PbPriceType.Limit, 16, 0)
    
    # 执行平仓
    QuickInsertOrder(context.myacc, g.code, 'sell', 'close', price_type, volume)
    
    # 计算盈亏
    profit = (price - g.entry_price) * volume
    is_win = profit > 0
    
    print(f"平多头仓位: {g.code}, 价格: {price}, 数量: {volume}, 盈亏: {profit}")
    
    # 更新交易统计
    g.trade_count += 1
    if is_win:
        g.win_count += 1
    
    # 记录交易
    g.trade_log.append({
        'time': GetCurrentTime(),
        'action': 'close_long',
        'code': g.code,
        'price': price,
        'volume': volume,
        'profit': profit,
        'is_win': is_win
    })
    
    # 重置持仓状态
    g.position = 0
    g.entry_price = 0

# 平空头仓位
def close_short_position(context, price):
    if g.position >= 0:
        return
    
    volume = abs(g.position)
    price_type = PriceType(PbPriceType.Limit, 16, 0)
    
    # 执行平仓
    QuickInsertOrder(context.myacc, g.code, 'buy', 'close', price_type, volume)
    
    # 计算盈亏
    profit = (g.entry_price - price) * volume
    is_win = profit > 0
    
    print(f"平空头仓位: {g.code}, 价格: {price}, 数量: {volume}, 盈亏: {profit}")
    
    # 更新交易统计
    g.trade_count += 1
    if is_win:
        g.win_count += 1
    
    # 记录交易
    g.trade_log.append({
        'time': GetCurrentTime(),
        'action': 'close_short',
        'code': g.code,
        'price': price,
        'volume': volume,
        'profit': profit,
        'is_win': is_win
    })
    
    # 重置持仓状态
    g.position = 0
    g.entry_price = 0

# 检查止盈止损
def check_stop_loss_take_profit(context, current_price):
    if g.position == 0:
        return
    
    # 计算价格变动百分比
    price_change_pct = abs(current_price - g.entry_price) / g.entry_price
    
    # 多头持仓
    if g.position > 0:
        # 止损
        if current_price < g.entry_price and price_change_pct >= g.stop_loss_pct:
            print(f"触发止损: 当前价格 {current_price}, 开仓价格 {g.entry_price}")
            close_long_position(context, current_price)
        
        # 止盈
        elif current_price > g.entry_price and price_change_pct >= g.take_profit_pct:
            print(f"触发止盈: 当前价格 {current_price}, 开仓价格 {g.entry_price}")
            close_long_position(context, current_price)
    
    # 空头持仓
    elif g.position < 0:
        # 止损
        if current_price > g.entry_price and price_change_pct >= g.stop_loss_pct:
            print(f"触发止损: 当前价格 {current_price}, 开仓价格 {g.entry_price}")
            close_short_position(context, current_price)
        
        # 止盈
        elif current_price < g.entry_price and price_change_pct >= g.take_profit_pct:
            print(f"触发止盈: 当前价格 {current_price}, 开仓价格 {g.entry_price}")
            close_short_position(context, current_price)

# 平掉所有持仓
def close_all_positions(context):
    if g.position > 0:
        close_long_position(context, GetQuote(g.code).now)
    elif g.position < 0:
        close_short_position(context, GetQuote(g.code).now)
    
    # 确保持仓状态重置
    g.position = 0
    g.entry_price = 0